import os
import json
from PIL import Image
from tqdm import tqdm
img_dir = os.path.join("E:\internship\dataset_test\cloud", "indi_doc\images")
dst_dir = os.path.join("E:\internship\dataset_test\cloud", "indi_doc\images_crop")

# ==================================================================================
os.makedirs(dst_dir,exist_ok=True)
ftxt = open(os.path.join(dst_dir, "Label.txt"), "w", encoding="utf-8")
count = 0
label_dir = img_dir.replace("images","labels")

fns = os.listdir(img_dir)
for fn in tqdm(fns):
    fgdt = open(os.path.join(label_dir, fn + ".txt"), "r", encoding="utf-8")
    fgdt_lines = fgdt.readlines()
    if len(fgdt_lines) > 1:
        print("!!!error : len(fgdt_lines)> 1")
    label_json = json.loads(fgdt_lines[0])
    if len(label_json) > 1:
        print("!!!error : len(label_json)> 1")
    path = os.path.join(img_dir, fn)
    path = path.replace("\\", "/")
    for current_label in label_json[0]["Labels"]:
        value = current_label['Value']
        if ('picture' not in value) and ('para' not in value) and ('Para' not in value) and ('table' not in value) and (
                value.split("-")[0] != 'cell'):
            xyxy = current_label['Points']
            if len(xyxy) == 4:
                point1 = [round(xyxy[0]["X"]), round(xyxy[0]["Y"])]
                point2 = [round(xyxy[1]["X"]), round(xyxy[1]["Y"])]
                point3 = [round(xyxy[2]["X"]), round(xyxy[2]["Y"])]
                point4 = [round(xyxy[3]["X"]), round(xyxy[3]["Y"])]
                xyxy = [point1, point2, point3, point4]
            if len(xyxy) == 2:
                point1 = [round(xyxy[0]["X"]), round(xyxy[0]["Y"])]
                point2 = [round(xyxy[1]["X"]), round(xyxy[0]["Y"])]
                point3 = [round(xyxy[1]["X"]), round(xyxy[1]["Y"])]
                point4 = [round(xyxy[0]["X"]), round(xyxy[1]["Y"])]
                xyxy = [point1, point2, point3, point4]
        x1 = 1e7
        y1 = 1e7
        x2 = 0
        y2 = 0
        for point in xyxy:
            if point[0] < x1:
                x1 = point[0]
            if point[0] > x2:
                x2 = point[0]
            if point[1] < y1:
                y1 = point[1]
            if point[1] > y2:
                y2 = point[1]
        im = Image.open(os.path.join(img_dir,fn))
        w, h = im.size
        cropim = im.crop((x1, y1 - 1, x2, y2 + 1))
        imgname = str(count) + ".png"
        cropim.save(os.path.join(dst_dir, imgname))
        value = value.replace("\n", "")
        ftxt.write(imgname + "\t" + value + "\n")
        count = count + 1
ftxt.close()